"""
PDF处理核心模块
"""

from typing import List, Tuple, Optional, Union
import fitz  # PyMuPDF
from PIL import Image
from PyPDF2 import PdfReader, PdfWriter
from pathlib import Path
import tempfile
import os

from ..utils.exceptions import PDFProcessingError, FileNotFoundError, InvalidFileFormatError
from ..utils.logger import Logger


class PDFProcessor:
    """PDF处理器"""
    
    def __init__(self, logger: Optional[Logger] = None):
        self.logger = logger or Logger("PDFProcessor")
        self.current_document: Optional[fitz.Document] = None
        self.current_file_path: Optional[str] = None
        self.page_count = 0
    
    def load_document(self, file_path: Union[str, Path]) -> bool:
        """
        加载PDF文档
        
        Args:
            file_path: PDF文件路径
            
        Returns:
            bool: 加载是否成功
            
        Raises:
            FileNotFoundError: 文件不存在
            InvalidFileFormatError: 文件格式无效
            PDFProcessingError: PDF处理错误
        """
        try:
            file_path = Path(file_path)
            
            if not file_path.exists():
                raise FileNotFoundError(f"PDF文件不存在: {file_path}")
            
            if not file_path.suffix.lower() == '.pdf':
                raise InvalidFileFormatError(f"不是有效的PDF文件: {file_path}")
            
            # 关闭之前的文档
            self.close_document()
            
            # 打开新文档
            self.current_document = fitz.open(str(file_path))
            self.current_file_path = str(file_path)
            self.page_count = len(self.current_document)
            
            self.logger.info(f"成功加载PDF文档: {file_path}, 页数: {self.page_count}")
            return True
            
        except Exception as e:
            self.logger.error(f"加载PDF文档失败: {e}", exc_info=True)
            if isinstance(e, (FileNotFoundError, InvalidFileFormatError)):
                raise
            raise PDFProcessingError(f"无法加载PDF文件: {e}")
    
    def close_document(self):
        """关闭当前文档"""
        if self.current_document:
            self.current_document.close()
            self.current_document = None
            self.current_file_path = None
            self.page_count = 0
    
    def generate_thumbnail(self, page_num: int, size: Tuple[int, int] = (220, 310)) -> Image.Image:
        """
        生成单页缩略图
        
        Args:
            page_num: 页码（从0开始）
            size: 缩略图尺寸 (width, height)
            
        Returns:
            PIL.Image: 缩略图图像
            
        Raises:
            PDFProcessingError: 生成缩略图失败
        """
        if not self.current_document:
            raise PDFProcessingError("未加载PDF文档")
        
        if page_num < 0 or page_num >= self.page_count:
            raise PDFProcessingError(f"页码超出范围: {page_num}, 总页数: {self.page_count}")
        
        try:
            page = self.current_document.load_page(page_num)
            
            # 计算缩放比例
            zoom = size[0] / page.rect.width
            mat = fitz.Matrix(zoom, zoom)
            
            # 生成像素图
            pix = page.get_pixmap(matrix=mat)
            
            # 转换为PIL图像
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            img.thumbnail(size, Image.Resampling.LANCZOS)
            
            return img
            
        except Exception as e:
            self.logger.error(f"生成缩略图失败 - 页码: {page_num}, 错误: {e}")
            raise PDFProcessingError(f"生成缩略图失败: {e}")
    
    def generate_all_thumbnails(self, size: Tuple[int, int] = (220, 310), 
                              progress_callback: Optional[callable] = None) -> List[Image.Image]:
        """
        生成所有页面的缩略图
        
        Args:
            size: 缩略图尺寸
            progress_callback: 进度回调函数 callback(current, total)
            
        Returns:
            List[Image.Image]: 缩略图列表
        """
        if not self.current_document:
            raise PDFProcessingError("未加载PDF文档")
        
        thumbnails = []
        
        for page_num in range(self.page_count):
            try:
                thumbnail = self.generate_thumbnail(page_num, size)
                thumbnails.append(thumbnail)
                
                if progress_callback:
                    progress_callback(page_num + 1, self.page_count)
                    
            except Exception as e:
                self.logger.warning(f"跳过页面 {page_num + 1}: {e}")
                # 创建空白缩略图作为占位符
                blank_img = Image.new('RGB', size, color='white')
                thumbnails.append(blank_img)
        
        return thumbnails

    def extract_pages(self, page_numbers: List[int], output_path: Union[str, Path]) -> bool:
        """
        提取指定页面到新PDF

        Args:
            page_numbers: 页码列表（从1开始）
            output_path: 输出文件路径

        Returns:
            bool: 提取是否成功
        """
        if not self.current_file_path:
            raise PDFProcessingError("未加载PDF文档")

        try:
            reader = PdfReader(self.current_file_path)
            writer = PdfWriter()

            # 添加选中的页面
            for page_num in sorted(page_numbers):
                if 1 <= page_num <= len(reader.pages):
                    writer.add_page(reader.pages[page_num - 1])
                else:
                    self.logger.warning(f"跳过无效页码: {page_num}")

            # 写入文件
            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)

            with open(output_path, "wb") as output_file:
                writer.write(output_file)

            self.logger.info(f"成功提取 {len(page_numbers)} 页到: {output_path}")
            return True

        except Exception as e:
            self.logger.error(f"提取页面失败: {e}", exc_info=True)
            raise PDFProcessingError(f"提取页面失败: {e}")

    def extract_pages_separately(self, page_numbers: List[int], output_dir: Union[str, Path],
                                filename_template: str = "{base_name}_page{page:03d}.pdf",
                                progress_callback: Optional[callable] = None) -> List[str]:
        """
        分别提取页面为单独的PDF文件

        Args:
            page_numbers: 页码列表（从1开始）
            output_dir: 输出目录
            filename_template: 文件名模板
            progress_callback: 进度回调函数

        Returns:
            List[str]: 成功创建的文件路径列表
        """
        if not self.current_file_path:
            raise PDFProcessingError("未加载PDF文档")

        try:
            reader = PdfReader(self.current_file_path)
            output_dir = Path(output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)

            base_name = Path(self.current_file_path).stem
            created_files = []

            for i, page_num in enumerate(sorted(page_numbers)):
                if 1 <= page_num <= len(reader.pages):
                    # 创建单页PDF
                    writer = PdfWriter()
                    writer.add_page(reader.pages[page_num - 1])

                    # 生成文件名
                    filename = filename_template.format(
                        base_name=base_name,
                        page=page_num,
                        index=i + 1
                    )
                    output_path = output_dir / filename

                    # 写入文件
                    with open(output_path, "wb") as output_file:
                        writer.write(output_file)

                    created_files.append(str(output_path))

                    if progress_callback:
                        progress_callback(i + 1, len(page_numbers))
                else:
                    self.logger.warning(f"跳过无效页码: {page_num}")

            self.logger.info(f"成功创建 {len(created_files)} 个PDF文件")
            return created_files

        except Exception as e:
            self.logger.error(f"分页提取失败: {e}", exc_info=True)
            raise PDFProcessingError(f"分页提取失败: {e}")

    def get_page_info(self, page_num: int) -> dict:
        """
        获取页面信息

        Args:
            page_num: 页码（从0开始）

        Returns:
            dict: 页面信息
        """
        if not self.current_document:
            raise PDFProcessingError("未加载PDF文档")

        if page_num < 0 or page_num >= self.page_count:
            raise PDFProcessingError(f"页码超出范围: {page_num}")

        try:
            page = self.current_document.load_page(page_num)
            rect = page.rect

            return {
                'page_number': page_num + 1,
                'width': rect.width,
                'height': rect.height,
                'rotation': page.rotation,
                'has_text': bool(page.get_text().strip()),
                'has_images': bool(page.get_images()),
            }

        except Exception as e:
            self.logger.error(f"获取页面信息失败: {e}")
            raise PDFProcessingError(f"获取页面信息失败: {e}")

    def __del__(self):
        """析构函数，确保文档被关闭"""
        self.close_document()
